/*
 * HSM Proxy Project.
 * Copyright (C) 2013 FedICT.
 * Copyright (C) 2011-2012 AGIV.
 *
 * This is free software; you can redistribute it and/or modify it
 * under the terms of the GNU Lesser General Public License version
 * 3.0 as published by the Free Software Foundation.
 *
 * This software is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
 * Lesser General Public License for more details.
 *
 * You should have received a copy of the GNU Lesser General Public
 * License along with this software; if not, see 
 * http://www.gnu.org/licenses/.
 */

package be.fedict.hsm.ws.impl;

import java.security.cert.X509Certificate;
import java.util.HashSet;
import java.util.List;
import java.util.Set;

import javax.ejb.EJB;
import javax.xml.namespace.QName;
import javax.xml.soap.SOAPBody;
import javax.xml.soap.SOAPEnvelope;
import javax.xml.soap.SOAPException;
import javax.xml.soap.SOAPMessage;
import javax.xml.soap.SOAPPart;
import javax.xml.ws.ProtocolException;
import javax.xml.ws.handler.MessageContext;
import javax.xml.ws.handler.soap.SOAPHandler;
import javax.xml.ws.handler.soap.SOAPMessageContext;

import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
import org.apache.ws.security.WSConstants;
import org.apache.ws.security.WSDataRef;
import org.apache.ws.security.WSSConfig;
import org.apache.ws.security.WSSecurityEngine;
import org.apache.ws.security.WSSecurityEngineResult;
import org.apache.ws.security.WSSecurityException;
import org.apache.ws.security.components.crypto.Crypto;
import org.apache.ws.security.message.token.BinarySecurity;
import org.apache.ws.security.message.token.Timestamp;
import org.apache.ws.security.util.WSSecurityUtil;

import be.fedict.hsm.model.security.SecurityAuditGeneratorBean;
import be.fedict.hsm.model.security.SecurityFunction;

/**
 * WS-Security verification SOAP handler. This SOAP handler checks whether a
 * WS-Security header is present. The WS-Security header should include a
 * timestamp and a signature over both the timestamp and the SOAP Body. The
 * received X509 certificate acts as the application credential.
 * <p/>
 * Even if we later on decide to use WS-SecurityPolicy via Apache CXF, we still
 * want to keep these additional checks in place.
 * 
 * @author Frank Cornelis
 * 
 */
@SecurityFunction("SF.I&A.1")
public class WSSecuritySOAPHandler implements SOAPHandler<SOAPMessageContext> {

	private static final Log LOG = LogFactory
			.getLog(WSSecuritySOAPHandler.class);

	private static final String X509_ATTRIBUTE = WSSecuritySOAPHandler.class
			.getName() + ".x509";

	private static final String WSU_NAMESPACE = "http://docs.oasis-open.org/wss/2004/01/oasis-200401-wss-wssecurity-utility-1.0.xsd";

	@EJB
	private SecurityAuditGeneratorBean securityAuditGeneratorBean;

	@Override
	public boolean handleMessage(SOAPMessageContext context) {
		LOG.debug("handleMessage");
		Boolean outboundProperty = (Boolean) context
				.get(MessageContext.MESSAGE_OUTBOUND_PROPERTY);
		if (false == outboundProperty.booleanValue()) {
			try {
				handleInboundMessage(context);
			} catch (Exception e) {
				throw new ProtocolException(e);
			}
		}
		return true;
	}

	private void handleInboundMessage(SOAPMessageContext context)
			throws WSSecurityException, SOAPException {
		LOG.debug("checking WS-Security header");
		SOAPMessage soapMessage = context.getMessage();
		SOAPPart soapPart = soapMessage.getSOAPPart();

		WSSecurityEngine secEngine = new WSSecurityEngine();
		Crypto crypto = new WSSecurityCrypto();
		WSSConfig wssConfig = new WSSConfig();
		wssConfig.setWsiBSPCompliant(true);
		secEngine.setWssConfig(wssConfig);
		List<WSSecurityEngineResult> results = secEngine.processSecurityHeader(
				soapPart, null, null, crypto);
		if (null == results) {
			this.securityAuditGeneratorBean.webServiceAuthenticationError();
			throw new SecurityException("no WS-Security results");
		}

		WSSecurityEngineResult timeStampActionResult = WSSecurityUtil
				.fetchActionResult(results, WSConstants.TS);
		if (null == timeStampActionResult) {
			this.securityAuditGeneratorBean.webServiceAuthenticationError();
			throw new SecurityException("no WS-Security timestamp result");
		}

		Timestamp receivedTimestamp = (Timestamp) timeStampActionResult
				.get(WSSecurityEngineResult.TAG_TIMESTAMP);
		if (null == receivedTimestamp) {
			this.securityAuditGeneratorBean.webServiceAuthenticationError();
			throw new SecurityException("no WS-Security timestamp");
		}

		LOG.debug("WS-Security timestamp created: "
				+ receivedTimestamp.getCreated());
		LOG.debug("WS-Security timestamp expires: "
				+ receivedTimestamp.getExpires());
		String timeStampIdRef = "#" + receivedTimestamp.getID();

		WSSecurityEngineResult bstActionResult = WSSecurityUtil
				.fetchActionResult(results, WSConstants.BST);
		if (null == bstActionResult) {
			this.securityAuditGeneratorBean.webServiceAuthenticationError();
			throw new SecurityException("no WS-Security BinarySecurityToken");
		}
		BinarySecurity binarySecurityToken = (BinarySecurity) bstActionResult
				.get(WSSecurityEngineResult.TAG_BINARY_SECURITY_TOKEN);

		WSSecurityEngineResult signActionResult = WSSecurityUtil
				.fetchActionResult(results, WSConstants.SIGN);
		if (null == signActionResult) {
			this.securityAuditGeneratorBean.webServiceAuthenticationError();
			throw new SecurityException("no valid XML signature");
		}
		String signatureMethod = (String) signActionResult
				.get(WSSecurityEngineResult.TAG_SIGNATURE_METHOD);
		LOG.debug("signature method: " + signatureMethod);
		if (false == "http://www.w3.org/2001/04/xmldsig-more#rsa-sha256"
				.equals(signatureMethod)) {
			this.securityAuditGeneratorBean.webServiceAuthenticationError();
			throw new SecurityException("signature algo should be RSA-SHA256");
		}
		X509Certificate certificate = (X509Certificate) signActionResult
				.get(WSSecurityEngineResult.TAG_X509_CERTIFICATE);
		LOG.debug("certificate subject: "
				+ certificate.getSubjectX500Principal());
		List<WSDataRef> wsDataRefs = (List<WSDataRef>) signActionResult
				.get(WSSecurityEngineResult.TAG_DATA_REF_URIS);

		SOAPEnvelope soapEnvelope = soapPart.getEnvelope();
		SOAPBody soapBody = soapEnvelope.getBody();
		String bodyIdRef = "#" + soapBody.getAttributeNS(WSU_NAMESPACE, "Id");
		String bstIdRef = "#" + binarySecurityToken.getID();

		boolean timestampDigested = false;
		boolean bodyDigested = false;
		boolean tokenDigested = false;
		for (WSDataRef wsDataRef : wsDataRefs) {
			String wsuId = wsDataRef.getWsuId();
			LOG.debug("signed wsu:Id: " + wsuId);
			LOG.debug("digest algorithm: " + wsDataRef.getDigestAlgorithm());
			if (false == "http://www.w3.org/2001/04/xmlenc#sha256"
					.equals(wsDataRef.getDigestAlgorithm())) {
				this.securityAuditGeneratorBean
						.webServiceAuthenticationError(certificate);
				throw new SecurityException("digest algorithm should be SHA256");
			}
			if (timeStampIdRef.equals(wsuId)) {
				timestampDigested = true;
			} else if (bodyIdRef.equals(wsuId)) {
				bodyDigested = true;
			} else if (bstIdRef.equals(wsuId)) {
				tokenDigested = true;
			}
		}
		if (false == timestampDigested) {
			this.securityAuditGeneratorBean
					.webServiceAuthenticationError(certificate);
			throw new SecurityException("timestamp not digested");
		}
		if (false == bodyDigested) {
			this.securityAuditGeneratorBean
					.webServiceAuthenticationError(certificate);
			throw new SecurityException("SOAP Body not digested");
		}
		if (false == tokenDigested) {
			this.securityAuditGeneratorBean
					.webServiceAuthenticationError(certificate);
			throw new SecurityException("BinarySecurityToken not digested");
		}

		context.put(X509_ATTRIBUTE, certificate);
	}

	@Override
	public boolean handleFault(SOAPMessageContext context) {
		return true;
	}

	@Override
	public void close(MessageContext context) {
	}

	@Override
	public Set<QName> getHeaders() {
		Set<QName> headers = new HashSet<QName>();
		headers.add(new QName(WSConstants.WSSE_NS, WSConstants.WSSE_LN));
		return headers;
	}

	public static X509Certificate getAuthenticatedCertificate(
			SOAPMessageContext context) {
		return (X509Certificate) context.get(X509_ATTRIBUTE);
	}
}
